package cn.sc.summer.mybatis.config.interceptor;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.ReflectUtil;
import cn.hutool.extra.spring.SpringUtil;
import cn.sc.summer.mybatis.context.DynamicDataSourceContext;
import com.baomidou.mybatisplus.core.conditions.AbstractWrapper;
import com.baomidou.mybatisplus.core.conditions.update.Update;
import com.baomidou.mybatisplus.core.toolkit.Constants;
import com.baomidou.mybatisplus.extension.plugins.inner.InnerInterceptor;
import cn.sc.summer.constant.encrypt.EncryptKeyConstant;
import cn.sc.summer.constant.mybatisplus.datasource.DataSourceConstant;
import cn.sc.summer.mybatis.annotation.encrypt.EncryptedField;
import cn.sc.summer.mybatis.annotation.encrypt.EncryptedTable;
import cn.sc.summer.mybatis.util.LoadBalancerUtil;
import cn.sc.summer.token.util.AESUtil;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.parser.CCJSqlParserManager;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.util.TablesNamesFinder;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.core.annotation.AnnotationUtils;

import javax.annotation.PostConstruct;
import java.io.StringReader;
import java.lang.reflect.Field;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * 类名：sql读写分离拦截器
 *
 * @author a-xin
 * @date 2024/5/29 15:32
 */
@Slf4j
public class DataSourceInterceptor implements InnerInterceptor, DisposableBean {

    private static Boolean READ_WRITE_ENABLE;

    public static final ThreadLocal<String> CURRENT_DB = new ThreadLocal<>();

    /**
     * 读写分离排除表信息
     */
    private static final String[] IGNORE_TABLE = {
            "sys_log",
            "gateway_log",
            "sms_log",
    };

    /**
     * 变量占位符正则
     */
    private static final Pattern PARAM_PAIRS_RE = Pattern.compile("#\\{ew\\.paramNameValuePairs\\.(" + Constants.WRAPPER_PARAM + "\\d+)\\}");

    @PostConstruct
    public static void init(Boolean enable) {
        DataSourceInterceptor.READ_WRITE_ENABLE = enable;
    }

    @Override
    public boolean willDoQuery(Executor executor, MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) throws SQLException {
        return InnerInterceptor.super.willDoQuery(executor, ms, parameter, rowBounds, resultHandler, boundSql);
    }

    /**
     * 读取数据
     */
    @Override
    public void beforeQuery(Executor executor, MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) {
        boolean needSetDatasource = needSetDatasource(boundSql);
        if (READ_WRITE_ENABLE && needSetDatasource) {
            CURRENT_DB.set(DynamicDataSourceContext.getDataSource());
            Map<String, List<String>> dataSourceMap = SpringUtil.getBean(DataSourceConstant.READ_WRITE_BEAN);
            List<String> readDataBases = dataSourceMap.get(DataSourceConstant.READ_DATABASE);
            if (readDataBases.size() > 1) {
                String nextServer = LoadBalancerUtil.getNextServer(readDataBases, DataSourceConstant.READ_PREFIX);
                DynamicDataSourceContext.setDataSource(nextServer);
                log.debug("==> READ ::: {}", nextServer);
            } else {
                DynamicDataSourceContext.setDataSource(readDataBases.get(0));
                log.debug("==> READ ::: {}", readDataBases.get(0));
            }
        }
    }

    @Override
    public boolean willDoUpdate(Executor executor, MappedStatement ms, Object parameter) throws SQLException {
        return InnerInterceptor.super.willDoUpdate(executor, ms, parameter);
    }

    /**
     * 操作数据，更新，写入，删除
     */
    @Override
    public void beforeUpdate(Executor executor, MappedStatement ms, Object parameter) {
        BoundSql boundSql = ms.getBoundSql(parameter);
        boolean needSetDatasource = needSetDatasource(boundSql);
        if (READ_WRITE_ENABLE && needSetDatasource) {
            CURRENT_DB.set(DynamicDataSourceContext.getDataSource());
            Map<String, List<String>> dataSourceMap = SpringUtil.getBean(DataSourceConstant.READ_WRITE_BEAN);
            List<String> readDataBases = dataSourceMap.get(DataSourceConstant.WRITE_DATABASE);
            if (readDataBases.size() > 1) {
                String nextServer = LoadBalancerUtil.getNextServer(readDataBases, DataSourceConstant.WRITE_PREFIX);
                DynamicDataSourceContext.setDataSource(nextServer);
                log.debug("==> WRITE ::: {}", nextServer);
            } else {
                DynamicDataSourceContext.setDataSource(readDataBases.get(0));
                log.debug("==> WRITE ::: {}", readDataBases.get(0));
            }
        }

        if (ObjectUtil.isNull(parameter)) {
            return;
        }

        // 通过MybatisPlus自带API（save、insert等）新增数据库时
        if (!(parameter instanceof Map)) {
            if (needToDecrypt(parameter.getClass())) {
                encryptEntity(parameter);
            }
            return;
        }

        Map<?, ?> paramMap = (Map<?, ?>) parameter;
        Object param;

        // 通过MybatisPlus自带API（update、updateById等）修改数据库时
        if (paramMap.containsKey(Constants.ENTITY) && null != (param = paramMap.get(Constants.ENTITY))) {
            if (needToDecrypt(param.getClass())) {
                encryptEntity(param);
            }
            return;
        }

        // 通过在mapper.xml中自定义API修改数据库时
        if (paramMap.containsKey("entity") && null != (param = paramMap.get("entity"))) {
            if (needToDecrypt(param.getClass())) {
                encryptEntity(param);
            }
            return;
        }

        // 通过UpdateWrapper、LambdaUpdateWrapper修改数据库时
        if (paramMap.containsKey(Constants.WRAPPER) && null != (param = paramMap.get(Constants.WRAPPER))) {
            if (param instanceof Update && param instanceof AbstractWrapper) {
                Class<?> entityClass = ms.getParameterMap().getType();
                if (needToDecrypt(entityClass)) {
                    encryptWrapper(entityClass, param);
                }
            }
        }

    }

    @Override
    public void beforePrepare(StatementHandler sh, Connection connection, Integer transactionTimeout) {
        InnerInterceptor.super.beforePrepare(sh, connection, transactionTimeout);
    }

    @Override
    public void beforeGetBoundSql(StatementHandler sh) {
        InnerInterceptor.super.beforeGetBoundSql(sh);
    }

    @Override
    public void setProperties(Properties properties) {
        InnerInterceptor.super.setProperties(properties);
    }

    /**
     * 校验该执行sql是否需要设置读写数据库信息
     *
     * @param boundSql 原始sql信息
     * @return 是否设置
     */
    private boolean needSetDatasource(BoundSql boundSql) {
        try {
            CCJSqlParserManager parserManager = new CCJSqlParserManager();
            Statement statement = parserManager.parse(new StringReader(boundSql.getSql()));
            TablesNamesFinder tablesNamesFinder = new TablesNamesFinder();
            List<String> tableNameList = tablesNamesFinder.getTableList(statement);
            if (!CollUtil.isEmpty(tableNameList)) {
                for (String table : IGNORE_TABLE) {
                    if (tableNameList.contains(table)) {
                        return false;
                    }
                }
            }
            return true;
        } catch (JSQLParserException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 校验该实例的类是否被@EncryptedTable所注解
     */
    private boolean needToDecrypt(Class<?> objectClass) {
        try {
            EncryptedTable sensitiveData = AnnotationUtils.findAnnotation(objectClass, EncryptedTable.class);
            return ObjectUtil.isNotNull(sensitiveData);
        } catch (Exception ex) {
            return false;
        }
    }

    /**
     * 通过API（save、updateById等）修改数据库时
     *
     * @param parameter 操作对象
     */
    private void encryptEntity(Object parameter) {
        //取出parameterType的类
        Class<?> resultClass = parameter.getClass();
        Field[] declaredFields = ReflectUtil.getFields(resultClass);
        for (Field field : declaredFields) {
            //取出所有被EncryptedField注解的字段
            EncryptedField encryptedField = field.getAnnotation(EncryptedField.class);
            if (ObjectUtil.isNotNull(encryptedField)) {
                field.setAccessible(true);
                try {
                    Object object = field.get(parameter);
                    if (object instanceof String) {
                        String value = (String) object;
                        field.set(parameter, AESUtil.encrypt(value, EncryptKeyConstant.DATASOURCE_AES_SECURE_KEY));
                    }
                } catch (IllegalAccessException ignored) {
                }
            }
        }
    }

    /**
     * 通过UpdateWrapper、LambdaUpdateWrapper修改数据库时
     *
     * @param entityClass c处理对象
     * @param ewParam     对象参数
     */
    private void encryptWrapper(Class<?> entityClass, Object ewParam) {
        AbstractWrapper<?, ?, ?> updateWrapper = (AbstractWrapper<?, ?, ?>) ewParam;
        String sqlSet = updateWrapper.getSqlSet();
        String[] elArr = sqlSet.split(",");
        Map<String, String> propMap = new HashMap<>(elArr.length);
        Arrays.stream(elArr).forEach(el -> {
            String[] elPart = el.split("=");
            propMap.put(elPart[0], elPart[1]);
        });

        //取出parameterType的类
        Field[] declaredFields = ReflectUtil.getFields(entityClass);
        for (Field field : declaredFields) {
            EncryptedField encryptedField = field.getAnnotation(EncryptedField.class);
            if (ObjectUtil.isNull(encryptedField)) {
                continue;
            }
            String el = propMap.get(field.getName());
            Matcher matcher = PARAM_PAIRS_RE.matcher(el);
            if (matcher.matches()) {
                String valueKey = matcher.group(1);
                Object value = updateWrapper.getParamNameValuePairs().get(valueKey);
                updateWrapper.getParamNameValuePairs().put(valueKey, AESUtil.encrypt(value.toString(), EncryptKeyConstant.DATASOURCE_AES_SECURE_KEY));
            }
        }
    }

    /**
     * Invoked by the containing {@code BeanFactory} on destruction of a bean.
     *
     */
    @Override
    public void destroy() {
        CURRENT_DB.remove();
    }

}
